import rediscluster
from config import br_redis_conf
from contextlib import contextmanager


class Rds:

    def __init__(self, redis_conf=br_redis_conf):
        self.redis_conf = redis_conf

    # 连接redis，非密码方式
    @contextmanager
    def connector(self, pipeline=False, decode=True):
        rc = rediscluster.RedisCluster(startup_nodes=self.redis_conf, decode_responses=decode)
        try:
            yield rc.pipeline() if pipeline else rc
        except UnicodeDecodeError:
            print("value重编码失败")
        finally:
            rc.close()

    # 执行查询
    def get(self, key_start: str, ids: list = None, value_type: str = "string"):
        if value_type.lower() not in ["string", None]:
            raise ValueError(f"value_type must be string or NONE")
        with self.connector(pipeline=True) as c:
            if ids:
                key = [f"{key_start}:{_id}" for _id in ids]
                [c.get(k) for k in key]
                return dict(zip(key, c.execute()))
            c.get(key_start)
            return {key_start: c.execute().pop()}

    def scan(self, key_start: str, value_type: str = "string"):
        if value_type.lower() not in ["string", None]:
            raise ValueError(f"value_type must be string or NONE")
        with self.connector() as c:
            cursor, count = 0, 0
            while count < 20:
                r = c.scan(cursor=cursor, match=key_start, count=1000)
                print(r)
                for _node, result in r.items():
                    cursor, r = result
                count += 1


rds = Rds()
r = rds.scan("TOUR:USER:MAIN:TASK")
print(r)
